package com.seehack.sample.thread.pool;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Java并发编程实战 8.4-扩展ThreadPoolExecutor
 * 对ThreadPoolExecutor进行扩展可以对执行线程进行统一拦截操作:
 * {@link ThreadPoolExecutor#beforeExecute(Thread, Runnable)}，
 * {@link ThreadPoolExecutor#afterExecute(Runnable, Throwable)}.
 */
public class TimeThreadPoolExecutor extends ThreadPoolExecutor {

    private static final ThreadLocal<Long> START_TIME_THREAD_LOCAL = ThreadLocal.withInitial(() -> 0L);
    private static final ThreadLocal<String> CURRENT_NAME_THREAD_LOCAL = ThreadLocal.withInitial(() -> Thread.currentThread().getName());
    private AtomicInteger counter = new AtomicInteger(0);

    public TimeThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
        setThreadFactory(r -> {
            /*
             * 设置线程池中work的名称
             * 实际线程由work执行,此处无法实现对每个实际执行线程进行名称修改
             */
            Thread thread = new Thread(r);
            thread.setName("TimeThreadPoolExecutor-" + counter.getAndIncrement());
            return thread;
        });
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        super.beforeExecute(t, r);
        START_TIME_THREAD_LOCAL.set(System.nanoTime());
    }

    /**
     * 该方法不论是否抛出异常都会执行
     *
     * @see ThreadPoolExecutor#runWorker(Worker)
     */
    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        try {
            long cost = System.nanoTime() - START_TIME_THREAD_LOCAL.get();
            System.out.println("name:" + CURRENT_NAME_THREAD_LOCAL.get() + ",cost time: " + cost + "ns");
        } finally {
            START_TIME_THREAD_LOCAL.remove();
            super.afterExecute(r, t);
        }
    }
}
